# !/usr/bin/python3
# -*- coding: utf-8 -*-
# ------------------------------------------
# @Time    : Date - 2021/8/21   Time - 15:09
# @Author  : Spence Guo Tang
# @FileName: models.py
# ------------------------------------------

import torch
import torch.nn as nn
import torch.nn.functional as F


# 使用卷积神经网络
class Classifier(nn.Module):
    def __init__(self):
        super(Classifier, self).__init__()
        self.conv_layer = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=(3, 3)),
            nn.MaxPool2d((3, 3)),
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=(3, 3)),
            nn.MaxPool2d((3, 3))
        )
        self.linear_layer = nn.Sequential(
            nn.Linear(in_features=128, out_features=64),
            nn.ReLU(),
            nn.Linear(in_features=64, out_features=10)
        )

    def forward(self, x):
        y = self.conv_layer(x)
        y = y.view(y.shape[0], -1)
        y = self.linear_layer(y)
        y = F.softmax(y, dim=-1)
        # y = torch.argmax(y, dim=-1)
        return y


# 基本全连接神经网络
class BaseModel(nn.Module):
    def __init__(self):
        super(BaseModel, self).__init__()
        self.fcn = nn.Sequential(
            nn.Linear(in_features=28*28, out_features=256),
            nn.PReLU(),
            nn.Linear(in_features=256, out_features=64),
            nn.PReLU(),
            nn.Linear(in_features=64, out_features=10)
        )

    def forward(self, x):
        x = x.view(x.shape[0], -1)
        y = self.fcn(x)
        outputs = torch.softmax(y, dim=1)
        outputs = torch.argmax(outputs, dim=1, keepdim=True)
        return outputs


if __name__ == '__main__':
    input = torch.randn((5, 1, 28, 28), dtype=torch.float)
    net = Classifier()
    output = net(input)
    print(output)
